# --------------------------------------------------------------
# @file         yrbasin_g09_predict.py
# @author       shenyc
# @description  使用模型进行预测
# --------------------------------------------------------------
import os
import time
import torch
import calendar
import matplotlib.pyplot as plt
import torch.nn.functional as F
from torch import Tensor
from datetime import date, datetime, timedelta, timezone
from typing import Iterable

import micaps
import yrbasin
import utils
# from yrbasin import CNNFeature12 as NET
# from yrbasin import CNNFeature13 as NET
from yrbasin import CNNFeature21 as NET
# --------------------------------------------------------------
from yrbasin_g09_common import (
    ROOT, TRANS_METHOD, PHYSCOUNT, OUTPUT_PATH,
    L1_FOLDER_RAIN, L1_FOLDER_STATION, L1_FOLDER_TRAIN, L1_FOLDER_PREDICT
)
# --------------------------------------------------------------
YEAR = 2022
# --------------------------------------------------------------
# 1 级
L1_FOLDER_PREDICT_YEAR = f"{L1_FOLDER_PREDICT}{YEAR}"   # 预测数据顶层文件夹
# --------------------------------------------------------------
# 2 级
L2_FOLDER_TRANS = f"T{TRANS_METHOD}"                # 版本号
# --------------------------------------------------------------
# 3 级
L3_FOLDER_FEATURE = f"Feature{PHYSCOUNT}"           # 特征世数据集文件夹
L3_FOLDER_TRAIN_OUTPUT = f"Output{PHYSCOUNT}"       # 训练输出文件夹
L3_FOLDER_PREDICT_OUTPUT = f"Predict{PHYSCOUNT}"    # 预测输出文件夹
# --------------------------------------------------------------
# 4 级
L4_FOLDER_TRAIN_MODEL = "Models"                    # 训练模型文件夹
# --------------------------------------------------------------
# 降水路径
RainPath = f"{ROOT}\\{L1_FOLDER_RAIN}"
# 站点路径
StationPath = f"{ROOT}\\{L1_FOLDER_STATION}"
# 特征路径
FeaturePath = f"{ROOT}\\{L1_FOLDER_PREDICT_YEAR}\\{L2_FOLDER_TRANS}\\{L3_FOLDER_FEATURE}"
# 模型路径
# ModelPath = f"{ROOT}\\{L1_FOLDER_TRAIN}\\{L2_FOLDER_TRANS}\\{L3_FOLDER_TRAIN_OUTPUT}\\km100\\{L4_FOLDER_TRAIN_MODEL}"
ModelPath = f"{ROOT}\\{L1_FOLDER_TRAIN}\\{L2_FOLDER_TRANS}\\{L3_FOLDER_TRAIN_OUTPUT}\\km200\\{L4_FOLDER_TRAIN_MODEL}"
# --------------------------------------------------------------
print(f"降水路径 RainPath = {RainPath}")
print(f"站点路径 StationPath = {StationPath}")
print(f"特征路径 FeaturePath = {FeaturePath}")
print(f"模型路径 ModelPath = {ModelPath}")
# --------------------------------------------------------------
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# --------------------------------------------------------------


class YrbStation(micaps.Station):
    """
    站点 - 黄河流域
    """

    def __init__(self):
        super(YrbStation, self).__init__()
        self.cnty: str = ""
        self.city: str = ""
        self._province: str = ""
        self.basin: str = ""
        self.basin_flag: int = 0
    # end __init__

    @property
    def province(self) -> str:
        return self._province

    @province.setter
    def province(self, val: str):
        """省"""
        self._province = val

# --------------------------------------------------------------


def get_net() -> NET:
    # file_name = "model_yrbasin_epoch_90.pth"
    file_name = "model_yrbasin_epoch_86.pth"
    full_name = os.path.join(ModelPath, file_name)
    return torch.load(full_name)


def get_feature_file(dt: datetime) -> str:
    file_name = "{0}20.bin".format(dt.strftime("%Y%m%d"))
    full_name = os.path.join(FeaturePath, file_name)
    return full_name


def get_station_infos() -> list[YrbStation]:
    file_name = "YrbStation.txt"
    full_name = os.path.join(StationPath, file_name)
    return read_yrb_station_info_(full_name)


def read_yrb_station_info_(fname: str) -> list[YrbStation]:
    """
    装载站点数据
    """
    alllines = ""
    with open(fname, 'r', encoding='utf-8') as infile:
        alllines = infile.readlines()
    # end with

    # list
    stations: list[YrbStation] = []
    for line in alllines:
        # xstation
        # 0,     1,      2,     3,     4,    5,     6,     7,       8, 9
        # 861002,108.299,36.780,1245.0,吴起县,延安市,陕西省,北洛河流域,2,吴起白豹
        xs = str(line).strip("\n").split(",")  # 去除首尾换行符，并按逗号划分
        station = YrbStation()
        station.id = xs[0]
        station.lon = float(xs[1])
        station.lat = float(xs[2])
        station.alt = float(xs[3])
        station.cnty = xs[4]
        station.city = xs[5]
        station.province = xs[6]
        station.basin = xs[7]
        station.basin_flag = int(xs[8])
        station.name = xs[9]
        # print(str(station))
        if (station.basin_flag != 1):
            stations.append(station)
        # endif
    # endfor
    return stations


def read_rain_info_(fname: str) -> list[YrbStation]:
    """
    装载雨量数据
    """
    alllines = ""
    with open(fname, 'r', encoding='utf-8') as infile:
        alllines = infile.readlines()
    # end with

    # list
    stations: list[YrbStation] = []
    for line in alllines:
        # xstation
        # 0,        1,      2,      3,      4
        # 861002	108.299	36.780	1245.0	6.3
        xs = str(line).replace(" ", "").strip("\n").split("\t")
        station = YrbStation()
        station.id = xs[0]
        station.lon = float(xs[1])
        station.lat = float(xs[2])
        station.alt = float(xs[3])
        station.value = float(xs[4])
        stations.append(station)
    # endfor
    # stations.remove()
    return stations

# -----------------------------------------------------------------------------
# -----------------------------------------------------------------------------


def predict(net: NET, feature_file: str, start_time: float) -> list[int]:
    """
    功能:
        预测

    参数:
        net -> NET: 网络模型
        feature_file ->str : 特征文件

    返回:
        list[int]: 预测列表
    """
    answers = []
    features = yrbasin.read_feature_file(feature_file)
    for i in range(features.size(0)):
        feature = features[i]
        feature = torch.Tensor(feature).unsqueeze(0)  # [12,9,9] => [1,12,9,9]
        answer = predict_one_feature(feature, net, DEVICE)
        answers.append(answer)
        # print(i, ":", 'anwser = ', answer)
    # end for
    t = time.time() - start_time
    print(f"{feature_file} -- OK! 耗时 {t}")
    return answers
# end predict


def predict_one_feature(feature: Tensor, net: NET, device: torch.device) -> int:
    """
    预测

    参数:
        feature:    特征数据样本
        net:        网络模型
        device:     torch.device

    返回:
        int: 预测值
    """
    net.to(device)
    with torch.no_grad():
        feature = feature.to(device)
        output: torch.Tensor = net(feature)     # 使用模型
        _, target = torch.max(output.data, 1)   # 每行最大值索引
        # 返回预测值, one of (0, 1, 2, 3, 4).
        return int(torch.Tensor(target).item())
    # end with
# end predict_on_feature


def output_micaps3(dt: datetime, stations: list[YrbStation], predicts: list[int]):
    dt_micaps = dt + timedelta(1)

    lines: list[str] = []
    for i in range(len(stations)):
        s = stations[i]
        v = predicts[i]
        line = "{0}\t{1:.3f}\t{2:.3f}\t{3:10.1f}\t{4:.1f}\n".format(
            s.id,
            s.lon,
            s.lat,
            s.alt,
            micaps.level_to_rain(v)
        )
        lines.append(line)
    # endfor

    # 插入 micaps 3 file header
    # diamond 3 2019年8月2日08时00分00秒黄河流域站点24小时降水
    # 2019 8 2 0 0
    # 0 0 0 0 1
    # 6734
    line = f"diamond 3 {dt_micaps.year}年{dt_micaps.month:02d}月{dt_micaps.day:02d}日24时00分00秒黄河流域站点24小时降水\n"
    lines.insert(0, line)

    line = f"{dt_micaps.year} {dt_micaps.month:02d} {dt_micaps.day:02d} 0 0\n"
    lines.insert(1, line)

    line = f"0 0 0 0 1\n"
    lines.insert(2, line)

    line = f"{len(stations)}\n"
    lines.insert(3, line)

    # 写文件
    micaps3_file = "{0}08.024".format(dt_micaps.strftime("%Y%m%d"))
    micaps3_file_path = os.path.join(OUTPUT_PATH, micaps3_file)
    with open(micaps3_file_path, 'w', encoding='utf-8') as f:
        f.writelines(lines)
    # endwith
    return


def predict_trans_method_xxxxx(dt_lst_lst: list[list[datetime]]):
    """
    predict_trans_method_xxxxx

    参数：
        dt_lst_lst: 日期组集合, 每月为一组。 [dts_202106, dts_202107, dts_202108, dts_202109]
    """
    start_time = time.time()
    # 网络
    net = get_net()
    # 站点信息
    stations = get_station_infos()

    for dt_lst in dt_lst_lst:
        for dt in dt_lst:
            #
            filePath = get_feature_file(dt)
            try:
                # 预测(每天预测一次)
                predicts = predict(net, filePath, start_time)
                # 预测输出
                output_micaps3(dt, stations, predicts)
                # 评分输出
                # output_grade(dt, stations, predicts)
                # f = get_feature_file(dt)
                # if (not os.path.exists(f)):
                #     print(f"{f} - 不存在")
            except FileNotFoundError:
                print(f'{filePath} -- File Not Found!')
            except PermissionError:
                print('PermissionError - 没有权限')
        # endfor
    # endfor
    t = time.time() - start_time
    print(f"总耗时 {t}")
    return


def predict_main(year: int):
    # 预测概要
    print(f"预测概要: {YEAR}年 - {PHYSCOUNT}个物理量 - {TRANS_METHOD} 号物理量转换方法")
    dts = []
    dts.append(utils.get_date_times(year, 6))
    dts.append(utils.get_date_times(year, 7))
    dts.append(utils.get_date_times(year, 8))
    dts.append(utils.get_date_times(year, 9))
    predict_trans_method_xxxxx(dts)
    pass


if __name__ == '__main__':
    x = 0
    predict_main(YEAR)
    pass


# ========================================================

def test(prefile: str):
    """
    评分

    参数：
        prefile: 预测雨量文件(micaps 3 文件)
    """
    # 雨量实况 (读取 R24)
    fname1 = rf"E:\sfxData\YRBASIN\R24\201806\2018060108.000"
    mf2 = micaps.MicapsFile(fname1)
    # 预测雨量
    fname2 = rf"E:\sfxData\YRBASIN\R24\201806\2018060108.000"
    mf2 = micaps.MicapsFile(fname1)
    # 输出对比
    # sf = "E:\\Temp\\{0}.csv".format(dt_out.strftime("%Y%m%d%H"))
    # output_sf(sf, r_stations, predicts)
    # print("{0} - ok!".format(dt.strftime("%Y%m%d")))
    pass
